/**
 * @file UI排序和重合批
 * @author DuskyHuang 圣子
*/

import { _decorator, Node, UITransform, RenderData, math, gfx, UI, UIRenderer, StencilManager } from 'cc';
import { JSB } from 'cc/env';
import { notifier } from '../../core/boot';
import { Notifications } from '../../configs';

let __ui_global_sorting__: boolean = true;

notifier.attach(Notifications.UI.ENABLE_UI_SORTING, () => __ui_global_sorting__ = true);
notifier.attach(Notifications.UI.DISABLE_UI_SORTING, () => __ui_global_sorting__ = false);

if (!('sortingPriority' in UITransform.prototype)) {
    Object.defineProperties(UITransform.prototype, {
        _sortingPriority: {
            configurable: true, enumerable: true,
            writable: true, value: 0
        },
        _sortingEnabled: {
            configurable: true, enumerable: true,
            writable: true, value: false
        }
    });

    Object.defineProperties(UITransform.prototype, {
        sortingPriority: {
            configurable: true, enumerable: true,
            get: function () { return this._sortingPriority ?? 0; },
            set: function (value) {
                this._sortingPriority = value;
                if (JSB) this.node.uiSortingPriority = value;
            }
        },
        sortingEnabled: {
            configurable: true, enumerable: true,
            get: function () { return !!this._sortingEnabled; },
            set: function (value) {
                this._sortingEnabled = value;
                if (JSB) this.node.uiSortingEnabled = value;
            }
        },
    });
}

function updateOpacity(renderData: RenderData, opacity: number) {
    const vfmt = renderData.vertexFormat;
    const vb = renderData.chunk.vb;
    let attr; let format; let stride;
    // Color component offset
    let offset = 0;
    for (let i = 0; i < vfmt.length; ++i) {
        attr = vfmt[i];
        format = gfx.FormatInfos[attr.format];
        if (format.hasAlpha) {
            stride = renderData.floatStride;
            if (format.size / format.count === 1) {
                const alpha = ~~math.clamp(Math.round(opacity * 255), 0, 255);
                // Uint color RGBA8
                for (let color = offset; color < vb.length; color += stride) {
                    vb[color] = ((vb[color] & 0xffffff00) | alpha) >>> 0;
                }
            } else if (format.size / format.count === 4) {
                // RGBA32 color, alpha at position 3
                for (let alpha = offset + 3; alpha < vb.length; alpha += stride) {
                    vb[alpha] = opacity;
                }
            }
        }
        offset += format.size >> 2;
    }
}

declare module 'cc' {

    interface UITransform {
        _sortingPriority: number;
        sortingPriority: number;
        _sortingEnabled: boolean;
        sortingEnabled: boolean;
    }

    interface UI {
        rendererCache: UIRenderer[];
        rendererOrder: boolean;
        flushRendererCache();
    }

    interface UIRenderer {
        renderPriority: number;
        renderOpacity: number;
    }
}

const T = UIRenderer;

Object.defineProperties(UI.prototype, {
    _rendererCache: { configurable: true, enumerable: false, writable: false, value: [] },
    flushRendererCache: {
        configurable: true, enumerable: false,
        value: function () {
            if (!__ui_global_sorting__) return;
            const cache = this._rendererCache;
            if (cache.length > 0) {
                if (this.rendererOrder) {
                    cache.sort((a, b) => a.renderPriority - b.renderPriority);
                }
                for (let render of cache) {
                    render.fillBuffers(this);
                    if (render.renderOpacity >= 0) {
                        updateOpacity(render.renderData, render.renderOpacity);
                        const buffer = render.renderData.getMeshBuffer();
                        if (buffer) {
                            buffer.setDirty();
                        }
                    }
                }
                cache.length = 0;
            }
            this.rendererOrder = false;
        }
    },
    update: {
        configurable: true, enumerable: false,
        value: function () {
            if (JSB) {
                return;
            }
            this.rendererOrder = false;
            const screens = this._screens;
            let offset = 0;
            for (let i = 0; i < screens.length; ++i) {
                const screen = screens[i];
                const scene = screen._getRenderScene();
                if (!screen.enabledInHierarchy || !scene) {
                    continue;
                }
                // Reset state and walk
                this._opacityDirty = 0;
                this._pOpacity = 1;

                this.walk(screen.node);
                if (__ui_global_sorting__) this.flushRendererCache();

                this.autoMergeBatches(this._currComponent!);
                this.resetRenderStates();

                let batchPriority = 0;
                if (this._batches.length > offset) {
                    for (; offset < this._batches.length; ++offset) {
                        const batch = this._batches.array[offset];

                        if (batch.model) {
                            const subModels = batch.model.subModels;
                            for (let j = 0; j < subModels.length; j++) {
                                subModels[j].priority = batchPriority++;
                            }
                        } else {
                            batch.descriptorSet = this._descriptorSetCache.getDescriptorSet(batch);
                        }
                        scene.addBatch(batch);
                    }
                }
            }
        }
    },
    walk: {
        configurable: true, enumerable: false,
        value: function (node: Node, level = 0, sortingPriority = 0) {
            if (!node.activeInHierarchy) {
                return;
            }
            const children = node.children;
            const uiProps = node._uiProps;
            const render = uiProps.uiComp as UIRenderer;
            const stencilEnterLevel = render && (render.stencilStage === 2 || render.stencilStage === 6);
            const transform = uiProps.uiTransformComp;
            if (__ui_global_sorting__) {
                sortingPriority = (transform && transform.sortingEnabled && transform.enabled) ? transform.sortingPriority : sortingPriority;
            }
            // Save opacity
            const parentOpacity = this._pOpacity;
            let opacity = parentOpacity;
            // TODO Always cascade ui property's local opacity before remove it
            const selfOpacity = render && render.color ? render.color.a / 255 : 1;
            this._pOpacity = opacity *= selfOpacity * uiProps.localOpacity;
            // TODO Set opacity to ui property's opacity before remove it
            // @ts-expect-error temporary force set, will be removed with ui property's opacity
            uiProps._opacity = opacity;
            if (uiProps.colorDirty) {
                // Cascade color dirty state
                this._opacityDirty++;
            }

            // Render assembler update logic
            if (render && render.enabledInHierarchy) {
                if (__ui_global_sorting__) {
                    if (stencilEnterLevel) {
                        this.flushRendererCache();

                        render.fillBuffers(this);// for rendering

                        // Update cascaded opacity to vertex buffer
                        if (this._opacityDirty && render && !render.useVertexOpacity && render.renderData && render.renderData.vertexCount > 0) {
                            // HARD COUPLING
                            updateOpacity(render.renderData, opacity);
                            const buffer = render.renderData.getMeshBuffer();
                            if (buffer) {
                                buffer.setDirty();
                            }
                        }
                    } else {
                        this._rendererCache.push(render);
                        render.renderPriority = sortingPriority;
                        if (sortingPriority != 0) {
                            this.rendererOrder = true;
                        }
                        if (this._opacityDirty && render && !render.useVertexOpacity && render.renderData && render.renderData.vertexCount > 0) {
                            render.renderOpacity = opacity;
                        } else {
                            render.renderOpacity = -1;
                        }
                    }
                } else render.fillBuffers(this);

            }

            if (children.length > 0 && !node._static) {
                for (let i = 0; i < children.length; ++i) {
                    const child = children[i];
                    this.walk(child, level, sortingPriority);
                }
            }

            if (uiProps.colorDirty) {
                // Reduce cascaded color dirty state
                this._opacityDirty--;
                // Reset color dirty
                uiProps.colorDirty = false;
            }
            // Restore opacity
            this._pOpacity = parentOpacity;
            // Post render assembler update logic
            // ATTENTION: Will also reset colorDirty inside postUpdateAssembler
            if (render && render.enabledInHierarchy) {
                render.postUpdateAssembler(this);
                if (stencilEnterLevel
                    && (StencilManager.sharedManager!.getMaskStackSize() > 0)) {
                    this.flushRendererCache();
                    this.autoMergeBatches(this._currComponent!);
                    this.resetRenderStates();
                    StencilManager.sharedManager!.exitMask();
                }
            }

            level += 1;
        }
    }
});